{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%%shell\n# Installs the latest dev build of TVM from PyPI. If you wish to build\n# from source, see https://tvm.apache.org/docs/install/from_source.html\npip install apache-tvm --pre"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n\n# 8. Creating Your MLPerfTiny Submission with microTVM\n**Authors**:\n[Mehrdad Hessar](https://github.com/mehrdadh)\n\nThis tutorial is showcasing building an MLPerfTiny submission using microTVM. This\ntutorial shows the steps to import a TFLite model from MLPerfTiny benchmark models,\ncompile it with TVM and generate a Zephyr project which can be flashed to a Zephyr\nsupported board to benchmark the model using EEMBC runner.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Install microTVM Python dependencies\n\nTVM does not include a package for Python serial communication, so\nwe must install one before using microTVM. We will also need TFLite\nto load models.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%%shell\npip install pyserial==3.5 tflite==2.1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import os\nimport pathlib\nimport tarfile\nimport tempfile\nimport shutil"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Install Zephyr\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%%shell\n# Install west and ninja\npython3 -m pip install west\napt-get install -y ninja-build\n\n# Install ZephyrProject\nZEPHYR_PROJECT_PATH=\"/content/zephyrproject\"\nexport ZEPHYR_BASE=${ZEPHYR_PROJECT_PATH}/zephyr\nwest init ${ZEPHYR_PROJECT_PATH}\ncd ${ZEPHYR_BASE}\ngit checkout v3.2-branch\ncd ..\nwest update\nwest zephyr-export\nchmod -R o+w ${ZEPHYR_PROJECT_PATH}\n\n# Install Zephyr SDK\ncd /content\nZEPHYR_SDK_VERSION=\"0.15.2\"\nwget \"https://github.com/zephyrproject-rtos/sdk-ng/releases/download/v${ZEPHYR_SDK_VERSION}/zephyr-sdk-${ZEPHYR_SDK_VERSION}_linux-x86_64.tar.gz\"\ntar xvf \"zephyr-sdk-${ZEPHYR_SDK_VERSION}_linux-x86_64.tar.gz\"\nmv \"zephyr-sdk-${ZEPHYR_SDK_VERSION}\" zephyr-sdk\nrm \"zephyr-sdk-${ZEPHYR_SDK_VERSION}_linux-x86_64.tar.gz\"\n\n# Install python dependencies\npython3 -m pip install -r \"${ZEPHYR_BASE}/scripts/requirements.txt\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "**Note:** Install CMSIS-NN only if you are interested to generate this submission\nusing CMSIS-NN code generator.\n\n\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Install CMSIS-NN\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%%shell\nCMSIS_SHA=\"51263182d16c92649a48144ba56c0945f9fce60e\"\nCMSIS_URL=\"http://github.com/ARM-software/CMSIS_5/archive/${CMSIS_SHA}.tar.gz\"\nexport CMSIS_PATH=/content/cmsis\nDOWNLOAD_PATH=\"/content/${CMSIS_SHA}.tar.gz\"\nmkdir ${CMSIS_PATH}\nwget ${CMSIS_URL} -O \"${DOWNLOAD_PATH}\"\ntar -xf \"${DOWNLOAD_PATH}\" -C ${CMSIS_PATH} --strip-components=1\nrm ${DOWNLOAD_PATH}\n\nCMSIS_NN_TAG=\"v4.0.0\"\nCMSIS_NN_URL=\"https://github.com/ARM-software/CMSIS-NN.git\"\ngit clone ${CMSIS_NN_URL} --branch ${CMSIS_NN_TAG} --single-branch ${CMSIS_PATH}/CMSIS-NN"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Import Python dependencies\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\nimport numpy as np\n\nimport tvm\nfrom tvm import relay\nfrom tvm.relay.backend import Executor, Runtime\nfrom tvm.contrib.download import download_testdata\nfrom tvm.micro import export_model_library_format\nimport tvm.micro.testing\nfrom tvm.micro.testing.utils import (\n    create_header_file,\n    mlf_extract_workspace_size_bytes,\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Import Visual Wake Word Model\n\nTo begin with, download and import the Visual Wake Word (VWW) TFLite model from MLPerfTiny.\nThis model is originally from [MLPerf Tiny repository](https://github.com/mlcommons/tiny).\nWe also capture metadata information from the TFLite model such as input/output name,\nquantization parameters, etc. which will be used in following steps.\n\nWe use indexing for various models to build the submission. The indices are defined as follows:\nTo build another model, you need to update the model URL, the short name and index number.\n\n  * Keyword Spotting(KWS) 1\n  * Visual Wake Word(VWW) 2\n  * Anomaly Detection(AD) 3\n  * Image Classification(IC) 4\n\nIf you would like to build the submission with CMSIS-NN, modify USE_CMSIS environment variable.\n\n```bash\nexport USE_CMSIS=1\n```\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "MODEL_URL = \"https://github.com/mlcommons/tiny/raw/bceb91c5ad2e2deb295547d81505721d3a87d578/benchmark/training/visual_wake_words/trained_models/vww_96_int8.tflite\"\nMODEL_PATH = download_testdata(MODEL_URL, \"vww_96_int8.tflite\", module=\"model\")\n\nMODEL_SHORT_NAME = \"VWW\"\nMODEL_INDEX = 2\n\nUSE_CMSIS = os.environ.get(\"TVM_USE_CMSIS\", False)\n\ntflite_model_buf = open(MODEL_PATH, \"rb\").read()\ntry:\n    import tflite\n\n    tflite_model = tflite.Model.GetRootAsModel(tflite_model_buf, 0)\nexcept AttributeError:\n    import tflite.Model\n\n    tflite_model = tflite.Model.Model.GetRootAsModel(tflite_model_buf, 0)\n\ninterpreter = tf.lite.Interpreter(model_path=str(MODEL_PATH))\ninterpreter.allocate_tensors()\ninput_details = interpreter.get_input_details()\noutput_details = interpreter.get_output_details()\n\ninput_name = input_details[0][\"name\"]\ninput_shape = tuple(input_details[0][\"shape\"])\ninput_dtype = np.dtype(input_details[0][\"dtype\"]).name\noutput_name = output_details[0][\"name\"]\noutput_shape = tuple(output_details[0][\"shape\"])\noutput_dtype = np.dtype(output_details[0][\"dtype\"]).name\n\n# We extract quantization information from TFLite model.\n# This is required for all models except Anomaly Detection,\n# because for other models we send quantized data to interpreter\n# from host, however, for AD model we send floating data and quantization\n# happens on the microcontroller.\nif MODEL_SHORT_NAME != \"AD\":\n    quant_output_scale = output_details[0][\"quantization_parameters\"][\"scales\"][0]\n    quant_output_zero_point = output_details[0][\"quantization_parameters\"][\"zero_points\"][0]\n\nrelay_mod, params = relay.frontend.from_tflite(\n    tflite_model, shape_dict={input_name: input_shape}, dtype_dict={input_name: input_dtype}\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Defining Target, Runtime and Executor\n\nNow we need to define the target, runtime and executor to compile this model. In this tutorial,\nwe use Ahead-of-Time (AoT) compilation and we build a standalone project. This is different\nthan using AoT with host-driven mode where the target would communicate with host using host-driven\nAoT executor to run inference.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Use the C runtime (crt)\nRUNTIME = Runtime(\"crt\")\n\n# Use the AoT executor with `unpacked-api=True` and `interface-api=c`. `interface-api=c` forces\n# the compiler to generate C type function APIs and `unpacked-api=True` forces the compiler\n# to generate minimal unpacked format inputs which reduces the stack memory usage on calling\n# inference layers of the model.\nEXECUTOR = Executor(\n    \"aot\",\n    {\"unpacked-api\": True, \"interface-api\": \"c\", \"workspace-byte-alignment\": 8},\n)\n\n# Select a Zephyr board\nBOARD = os.getenv(\"TVM_MICRO_BOARD\", default=\"nucleo_l4r5zi\")\n\n# Get the the full target description using the BOARD\nTARGET = tvm.micro.testing.get_target(\"zephyr\", BOARD)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Compile the model and export model library format\n\nNow, we compile the model for the target. Then, we generate model\nlibrary format for the compiled model. We also need to calculate the\nworkspace size that is required for the compiled model.\n\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "config = {\"tir.disable_vectorize\": True}\nif USE_CMSIS:\n    from tvm.relay.op.contrib import cmsisnn\n\n    config[\"relay.ext.cmsisnn.options\"] = {\"mcpu\": TARGET.mcpu}\n    relay_mod = cmsisnn.partition_for_cmsisnn(relay_mod, params, mcpu=TARGET.mcpu)\n\nwith tvm.transform.PassContext(opt_level=3, config=config):\n    module = tvm.relay.build(\n        relay_mod, target=TARGET, params=params, runtime=RUNTIME, executor=EXECUTOR\n    )\n\ntemp_dir = tvm.contrib.utils.tempdir()\nmodel_tar_path = temp_dir / \"model.tar\"\nexport_model_library_format(module, model_tar_path)\nworkspace_size = mlf_extract_workspace_size_bytes(model_tar_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Generate input/output header files\n\nTo create a microTVM standalone project with AoT, we need to generate\ninput and output header files. These header files are used to connect\nthe input and output API from generated code to the rest of the\nstandalone project. For this specific submission, we only need to generate\noutput header file since the input API call is handled differently.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "extra_tar_dir = tvm.contrib.utils.tempdir()\nextra_tar_file = extra_tar_dir / \"extra.tar\"\n\nwith tarfile.open(extra_tar_file, \"w:gz\") as tf:\n    create_header_file(\n        \"output_data\",\n        np.zeros(\n            shape=output_shape,\n            dtype=output_dtype,\n        ),\n        \"include/tvm\",\n        tf,\n    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Create the project, build and prepare the project tar file\n\nNow that we have the compiled model as a model library format,\nwe can generate the full project using Zephyr template project. First,\nwe prepare the project options, then build the project. Finally, we\ncleanup the temporary files and move the submission project to the\ncurrent working directory which could be downloaded and used on\nyour development kit.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "input_total_size = 1\nfor i in range(len(input_shape)):\n    input_total_size *= input_shape[i]\n\ntemplate_project_path = pathlib.Path(tvm.micro.get_microtvm_template_projects(\"zephyr\"))\nproject_options = {\n    \"extra_files_tar\": str(extra_tar_file),\n    \"project_type\": \"mlperftiny\",\n    \"board\": BOARD,\n    \"compile_definitions\": [\n        f\"-DWORKSPACE_SIZE={workspace_size + 512}\",  # Memory workspace size, 512 is a temporary offset\n        # since the memory calculation is not accurate.\n        f\"-DTARGET_MODEL={MODEL_INDEX}\",  # Sets the model index for project compilation.\n        f\"-DTH_MODEL_VERSION=EE_MODEL_VERSION_{MODEL_SHORT_NAME}01\",  # Sets model version. This is required by MLPerfTiny API.\n        f\"-DMAX_DB_INPUT_SIZE={input_total_size}\",  # Max size of the input data array.\n    ],\n}\n\nif MODEL_SHORT_NAME != \"AD\":\n    project_options[\"compile_definitions\"].append(f\"-DOUT_QUANT_SCALE={quant_output_scale}\")\n    project_options[\"compile_definitions\"].append(f\"-DOUT_QUANT_ZERO={quant_output_zero_point}\")\n\nif USE_CMSIS:\n    project_options[\"compile_definitions\"].append(f\"-DCOMPILE_WITH_CMSISNN=1\")\n\n# Note: You might need to adjust this based on the board that you are using.\nproject_options[\"config_main_stack_size\"] = 4000\n\nif USE_CMSIS:\n    project_options[\"cmsis_path\"] = os.environ.get(\"CMSIS_PATH\", \"/content/cmsis\")\n\ngenerated_project_dir = temp_dir / \"project\"\n\nproject = tvm.micro.project.generate_project_from_mlf(\n    template_project_path, generated_project_dir, model_tar_path, project_options\n)\nproject.build()\n\n# Cleanup the build directory and extra artifacts\nshutil.rmtree(generated_project_dir / \"build\")\n(generated_project_dir / \"model.tar\").unlink()\n\nproject_tar_path = pathlib.Path(os.getcwd()) / \"project.tar\"\nwith tarfile.open(project_tar_path, \"w:tar\") as tar:\n    tar.add(generated_project_dir, arcname=os.path.basename(\"project\"))\n\nprint(f\"The generated project is located here: {project_tar_path}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Use this project with your board\n\nNow that we have the generated project, you can use this project locally\nto flash your board and prepare it for EEMBC runner software.\nTo do this follow these steps:\n\n```bash\ntar -xf project.tar\ncd project\nmkdir build\ncmake ..\nmake -j2\nwest flash\n```\nNow you can connect your board to EEMBC runner using this\n[instructions](https://github.com/eembc/energyrunner)\nand benchmark this model on your board.\n\n\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.16"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}